import pandas as pd
from fbprophet import Prophet
from fbprophet.plot import add_changepoints_to_plot
import numpy as np
import matplotlib.pyplot as plt

# 读入数据集
live_df = pd.read_csv('data/data3.csv')
print(live_df.head())

# 定义节假日
chinese_holiday = pd.DataFrame({
    'holiday': 'Lunar_festivals',
    'ds': pd.to_datetime([
        '2024-02-10', '2024-02-11', '2024-02-12', '2024-02-13', '2024-02-14', '2024-02-15', '2024-02-16', '2024-02-17',
        '2024-04-04', '2024-04-05', '2024-04-06',
        '2024-06-08', '2024-06-09', '2024-06-10',
        '2024-09-15', '2024-09-16', '2024-09-17'
        ,
        '2025-01-26', '2025-01-27', '2025-01-28', '2025-01-29', '2025-01-30', '2025-01-31', '2025-02-01', '2025-02-02', '2025-02-03', '2025-02-04',
        '2025-04-04', '2025-04-05', '2025-04-06',
        '2025-05-31', '2025-06-01', '2025-06-02',
        '2025-10-06'
    ]),
    'lower_window': 0,
    'upper_window': 1,
})
china_holiday = pd.DataFrame({
    'holiday': 'china',
    'ds': pd.to_datetime([
        '2024-01-01',
        '2024-05-01', '2024-05-02', '2024-05-03', '2024-05-04', '2024-05-05',
        '2024-10-01', '2024-10-02', '2024-10-03', '2024-10-04', '2024-10-05', '2024-10-06', '2024-10-07'
        ,
        '2025-01-01',
        '2025-05-01', '2025-05-02', '2025-05-03', '2025-05-04', '2025-05-05',
        '2025-10-01', '2025-10-02', '2025-10-03', '2025-10-04', '2025-10-05', '2025-10-06', '2025-10-07', '2025-10-08'
    ]),
    'lower_window': 0,
    'upper_window': 1,
})
holidays = pd.concat((chinese_holiday, china_holiday))

model = Prophet(holidays=holidays)
model.fit(live_df)

future = model.make_future_dataframe(periods=4, freq='w', include_history=False)
future.tail(4)

forecast = model.predict(future)
forecast[['ds', 'yhat', 'yhat_lower', 'yhat_upper']].tail(4)

fig1 = model.plot(forecast)

fig = model.plot_components(forecast)

plt.show()
